import os
import sys

CC = 'mpicc'
CXX = 'mpicxx'
CFLAGS = "'-fopenmp -D_REENTRANT -msse4.1 -g -O3 -std=gnu99 -fPIC'"
CXXFLAGS = "'-fopenmp -D_REENTRANT -msse4.1 -g -O3 -std=c++17 -fexceptions -frtti -fPIC'"
#prefix = "/home/guilongcheng/chroma/lj_chroma_qdpjit_quda10/build"
build = "./build"
prefix = "/vopt"
nmake = 12
cuda_toolkit="/usr/local/cuda"

# -fpermissive may be needed when set --enable-alignment=64 for qdpxx
qmp_src="/src/qmp/"
llvm_13_src="/src/llvm-13.0.1.src/"
qdpxx_src="/src/qdpxx/"
qdp_jit_src="/src/qdp-jit/"
quda_src="/src/quda/"
chroma_src="/src/chroma/"
if len(sys.argv) >= 3:
	sm=sys.argv[2]
else:
	sm = "sm_60"
if_qmp,if_qdpxx,if_llvm_13,if_qdp_jit,if_quda,if_chroma,if_chroma_jit = [False]*7
if sys.argv[1]=="qmp":
	if_qmp = True
if sys.argv[1]=="qdpxx":
	if_qdpxx = True
if sys.argv[1]=="llvm":
	if_llvm_13 = True
if sys.argv[1]=="qdp_jit":
	if_qdp_jit = True
if sys.argv[1]=="quda":
	if_quda = True
if sys.argv[1]=="chroma":
	if_chroma_jit = True
if sys.argv[1] == "all":
	if_qmp,if_qdpxx,if_llvm_13,if_qdp_jit,if_quda,if_chroma_jit = [False]*6
#if_qmp = True
#if_qdpxx = True
#if_llvm_13 = True
#if_qdp_jit = True
#if_quda = True
#if_chroma = False
#if_chroma_jit = True

#########################################################################################################################

if if_qmp:
#	try:
#		os.system('git clone --recursive https://github.com/usqcd-software/qmp.git')
#	except:
#		pass
#
#	os.chdir('qmp')
#	os.system('pwd')
#	os.system('git checkout devel')
	os.chdir(prefix)
	os.system("mkdir -p qmp")
	os.chdir('qmp')
	os.system('pwd')

	os.chdir(qmp_src)
	os.system('git checkout devel')
	os.system("mkdir -p build")
	os.chdir('build')


	os.environ["CC"] = CC
	os.environ["CXX"] = CXX
	os.environ["CFLAGS"] = CFLAGS[1:len(CFLAGS)-1]
	os.environ["CXXFLAGS"] = CXXFLAGS[1:len(CXXFLAGS)-1]

	os.system('cmake ' +
	'-DCMAKE_INSTALL_PREFIX=%s/qmp '%prefix +
	'-DBUILD_TESTING=OFF ' +
	'-DQMP_TESTING=OFF ' +
	'-DQMP_MPI=ON %s'%qmp_src)
	os.system('make -j %d'%nmake)
	os.system('make install')
	os.chdir('/src')
	os.system('pwd')

	print('qmp done')

#########################################################################################################################

if if_qdpxx:
#	try:
#		os.system('git clone --recursive https://github.com/usqcd-software/qdpxx.git')
#	except:
#		pass

#	os.chdir('qdpxx')
	os.chdir(qdpxx_src)
	os.system('pwd')
	os.system('git checkout devel')

	os.chdir(prefix)
	os.system("mkdir -p qdpxx")
	os.chdir('qdpxx')
	os.system('pwd')

	os.environ["CC"] = CC
	os.environ["CXX"] = CXX
	os.environ["CFLAGS"] = CFLAGS[1:len(CFLAGS)-1]
	os.environ["CXXFLAGS"] = CXXFLAGS[1:len(CXXFLAGS)-1]

	os.system('cmake ' +
	'-DCMAKE_INSTALL_PREFIX=%s/qdpxx '%prefix +
	'-DQMP_DIR=%s/qmp/lib/cmake/QMP '%prefix +
	'-DQDP_USE_SSE=ON '+
	'-DQDP_PARALLEL_ARCH=parscalar '+
	'-DQDP_USE_SSE2=ON '+
	'-DQDP_USE_SSE3=ON '+
	'-DQDP_USE_OPENMP=ON %s'%qdpxx_src)
	os.system('make -j %d'%nmake)
	os.system('make install')
	os.chdir(prefix)
	os.chdir('..')
	os.system('pwd')

	print('qdpxx done')

#########################################################################################################################

if if_llvm_13:
	os.chdir(prefix)
	os.system("mkdir -p llvm-13")

	os.chdir(llvm_13_src)
	os.system("mkdir -p build")
	os.chdir('build')
	os.environ["CC"] = CC
	os.environ["CXX"] = CXX
	os.environ["CFLAGS"] = CFLAGS[1:len(CFLAGS)-1]
	os.environ["CXXFLAGS"] = CXXFLAGS[1:len(CXXFLAGS)-1]

	os.system('cmake ' +
	'-DLLVM_ENABLE_PROJECTS="llvm" '+
	'-DLLVM_ENABLE_TERMINFO="OFF" '+
	'-DCMAKE_BUILD_TYPE=Release '+
	'-DCMAKE_INSTALL_PREFIX=%s/llvm-13 '%prefix+
	'-DLLVM_TARGETS_TO_BUILD="X86;NVPTX" '+
	'-DLLVM_ENABLE_ZLIB="OFF" '+
	'-DBUILD_SHARED_LIBS="ON" '+
	'-DLLVM_ENABLE_RTTI="ON"  %s '%llvm_13_src)
	os.system('make -j %d'%nmake)
	os.system('make install')
	os.chdir('/src')
	os.system('pwd')

	print('llvm-13 done')

#########################################################################################################################
if if_qdp_jit:
#	try:
#		os.system('git clone --recursive https://github.com/JeffersonLab/qdp-jit')
#	except:
#		pass
#
#	os.chdir('qdp-jit')
#	os.system('pwd')
	os.chdir(prefix)
	os.system("mkdir -p qdp-jit")

	os.chdir(qdp_jit_src)
	os.system('git checkout devel')
	os.system("mkdir -p build")
	os.chdir('build')

	os.environ["CC"] = CC
	os.environ["CXX"] = CXX
	os.environ["CFLAGS"] = CFLAGS[1:len(CFLAGS)-1]
	os.environ["CXXFLAGS"] = CXXFLAGS[1:len(CXXFLAGS)-1] + ' -I%s/llvm-13/include '%prefix

	os.system('cmake ' +
	'-DCMAKE_INSTALL_PREFIX=%s/qdp-jit '%prefix +
	'-DQMP_DIR=%s/qmp/lib/cmake/QMP '%prefix +
	'-DCUDA_TOOLKIT_ROOT_DIR=%s '%cuda_toolkit+
	'-DLIME_LARGEGFILE=ON '+
	'-DQIO_ENABLE_PARALLEL_IO=ON '+
	'-DQIO_ENABLE_QMP_ROUTE=ON '+
	'-DBUILD_TESTING=OFF ' +
	'-DQDP_BUILD_EXAMPLES=OFF '+
	'-DQDP_ENABLE_PROFILING=OFF '+
	'-DQDP_ENABLE_CUSTOM_KERNELS=OFF '+   # why? any effects?
	'-DQDP_ENABLE_LLVM14=OFF '+
	'-DLLVM_DIR=%s/llvm-13/cmake/modules/CMakeFiles '%prefix +
	'-DQDP_USE_OPENMP=ON %s'%qdp_jit_src)
	os.system('make -j %d'%nmake)
	os.system('make install')
	os.chdir(prefix)
	os.chdir('..')
	os.system('pwd')

	print('qdp-jit done')
#########################################################################################################################

if if_quda:
#	try:
#		os.system('git clone --recursive git@github.com:Lattice/quda.git')
#	except:
#		pass

	os.chdir(prefix)
	os.system("mkdir -p quda_%s"%sm)

	os.chdir(quda_src)
	os.system('pwd')
	os.system('git checkout develop')
	os.system("mkdir -p build")
	os.chdir('build')

	os.environ["CC"] = CC
	os.environ["CXX"] = CXX
	os.environ["CFLAGS"] = CFLAGS[1:len(CFLAGS)-1]
	os.environ["CXXFLAGS"] = CXXFLAGS[1:len(CXXFLAGS)-1]
	os.environ["CUDA_INSTALL_PATH"] = cuda_toolkit
	os.environ["PK_CUDA_HOME"] = cuda_toolkit
	os.environ["PK_GPU_ARCH"] = sm
	if sm == "sm_35":
		print("ATOMIC is off!")
		ATOMIC="-DQUDA_HETEROGENEOUS_ATOMIC=OFF "
	else:
		print("ATOMIC is on!")
		ATOMIC=""

	os.system('cmake -Wno-dev ' +
	'-DMPI_CXX_COMPILER=mpicxx '+
	'-DMPI_C_COMPILER=mpicc '+
	'-DQUDA_DIRAC_CLOVER=ON '+
	'-DQUDA_DIRAC_DOMAIN_WALL=OFF '+
	'-DQUDA_DIRAC_NDEG_TWISTED_MASS=OFF '+
	'-DQUDA_DIRAC_STAGGERED=OFF '+
	'-DQUDA_DIRAC_TWISTED_MASS=ON '+
	'-DQUDA_DIRAC_TWISTED_CLOVER=ON '+
	'-DQUDA_DIRAC_WILSON=ON '+
	'-DQUDA_GPU_ARCH=%s '%sm+
	ATOMIC+ 
	'-DQUDA_INTERFACE_QDPJIT=ON '+
	'-DQUDA_QDPJIT=ON '+
	'-DQUDA_INTERFACE_MILC=OFF '+
	'-DQUDA_INTERFACE_CPS=OFF '+
	'-DQUDA_INTERFACE_QDP=ON '+
	'-DQUDA_INTERFACE_TIFR=OFF '+
	'-DQUDA_MPI=OFF '+
	'-DQUDA_QMP=ON '+
	'-DQDPXX_DIR=%s/qdp-jit/lib/cmake/QDPXX '%prefix +
	'-DLLVM_DIR=%s/llvm-13/cmake/modules/CMakeFiles '%prefix +
	'-DQMP_DIR=%s/qmp/lib/cmake/QMP '%prefix +
	'-DQIO_DIR=%s/qdpxx/lib/cmake/QIO '%prefix +
	'-DCLime_DIR=%s/qdpxx/lib/cmake/CLime '%prefix +
	'-DQUDA_QIO=ON '+
	'-DQUDA_MULTIGRID=ON '+
	'-DQUDA_MAX_MULTI_BLAS_N=9 '+
	'-DQUDA_DOWNLOAD_EIGEN=YES '+
	'-DCMAKE_INSTALL_PREFIX=%s/quda_%s '%(prefix,sm) +
	'-DCMAKE_BUILD_TYPE=RELEASE '+
#	'-DQUDA_BUILD_SHAREDLIB=NO '+
	'-DCUDA_cuda_driver_LIBRARY=%s/lib64/stubs/libcuda.so '%cuda_toolkit+
	'-DQUDA_BUILD_ALL_TESTS=OFF %s'%quda_src)
	os.system('make -j %d'%nmake)
	os.system('make install')
	os.chdir('/src')
	os.system('pwd')

	print('quda_%s done'%sm)

#########################################################################################################################
if if_chroma:
	try:
		os.system('git clone --recursive https://github.com/JeffersonLab/chroma.git')
	except:
		pass

	os.chdir(prefix)
	os.system("mkdir -p chroma")

	os.chdir(chroma_src)
	os.system('pwd')
	os.system('git checkout devel')
	os.system("mkdir -p build")
	os.chdir('build')

	os.environ["CC"] = CC
	os.environ["CXX"] = CXX
	os.environ["CFLAGS"] = CFLAGS[1:len(CFLAGS)-1]
	os.environ["CXXFLAGS"] = CXXFLAGS[1:len(CXXFLAGS)-1]

	os.system('cmake ' +
	'-DCMAKE_INSTALL_PREFIX=%s/chroma '%prefix +
	'-DChroma_ENABLE_LAPACK=ON '+
	'-DChroma_ENABLE_JIT_CLOVER=OFF '+
	'-DChroma_ENABLE_OPENMP=ON '+
	'-DChroma_ENABLE_MGPROTO=ON '+
	'-DChroma_ENABLE_QUDA=OFF '+
	'-DChroma_ENABLE_SSE2=ON '+
	'-DChroma_ENABLE_SSE3=ON '+
	'-DChroma_USE_MGPROTO=ON '+
	'-DQDPXX_DIR=%s/qdpxx/lib/cmake/QDPXX '%prefix +
	'-DQMP_DIR=%s/qmp/lib/cmake/QMP '%prefix +
	' %s'%chroma_src)

	os.system('make -j %d'%nmake)
	os.system('make install')
	os.chdir(prefix)
	os.chdir('..')
	os.system('pwd')

	print('chroma done')

#########################################################################################################################


if if_chroma_jit:
#	try:
#		os.system('git clone --recursive https://github.com/JeffersonLab/chroma.git')
#	except:
#		pass

#	os.chdir('chroma')
	os.chdir(prefix)
	os.system("mkdir -p chroma-jit")
	os.chdir('chroma-jit')
	os.system('pwd')

	os.chdir(chroma_src)
	os.system('pwd')
	os.system('git checkout devel')
	os.system("mkdir -p build")
	os.chdir('build')


	os.environ["CC"] = CC
	os.environ["CXX"] = CXX
	os.environ["CFLAGS"] = CFLAGS[1:len(CFLAGS)-1]
	os.environ["CXXFLAGS"] = CXXFLAGS[1:len(CXXFLAGS)-1]

	os.system('cmake ' +
	'-DCMAKE_INSTALL_PREFIX=%s/chroma-jit '%prefix +
	'-DChroma_ENABLE_LAPACK=ON '+
	'-DChroma_ENABLE_JIT_CLOVER=ON '+
	'-DChroma_ENABLE_OPENMP=ON '+
	'-DChroma_ENABLE_MGPROTO=ON '+
	'-DChroma_ENABLE_QUDA=ON '+
	'-DChroma_ENABLE_SSE2=ON '+
	'-DChroma_ENABLE_SSE3=ON '+
	'-DChroma_USE_MGPROTO=ON '+
	'-DQDPXX_DIR=%s/qdp-jit/lib/cmake/QDPXX '%prefix +
	'-DLLVM_DIR=%s/llvm-13/cmake/modules/CMakeFiles '%prefix +
	'-DQMP_DIR=%s/qmp/lib/cmake/QMP '%prefix +
	'-DQUDA_DIR=%s/quda_%s/lib/cmake/QUDA '%(prefix,sm) +
	' %s'%chroma_src)

	os.system('make -j %d'%nmake)
	os.system('make install')
	os.chdir('/src')
	os.system('pwd')

	print('chroma-jit done')

#########################################################################################################################
exit(0)
